package org.example.smartEducationBackend.controller;

import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatModel;
import jakarta.servlet.http.HttpServletResponse;
import lombok.RequiredArgsConstructor;
import org.example.smartEducationBackend.entity.POJO.RestResult;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.moderation.Generation;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.net.http.HttpResponse;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

@RestController
@RequestMapping("/AI")
@CrossOrigin(origins = "*")
@RequiredArgsConstructor
public class AIController {
    private final DashScopeChatModel chatModel;

    private final ExecutorService executor = Executors.newSingleThreadExecutor();

    @GetMapping(value = "/chat", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> chat(@RequestParam String prompt, HttpServletResponse response) {
        Flux<ChatResponse> stream = chatModel.stream(new Prompt(prompt));

        return stream.map(resp -> resp.getResult().getOutput().getContent())
                .doOnComplete(() -> System.out.println("Stream completed")); // 确保流正确完成
    }

    @GetMapping(value = "/sse/chat", produces = "text/event-stream")
    public SseEmitter chat(@RequestParam String prompt) {
        SseEmitter sseEmitter = new SseEmitter(0L);
        // 使用线程池异步处理Flux流
        executor.execute(() -> {
            try {
                // 将Flux流中的每个响应发送到SseEmitter
                chatModel.stream(new Prompt(prompt))
                        .map(resp -> resp.getResult().getOutput().getContent())
                        .subscribe(
                                content -> {
                                    try {
                                        sseEmitter.send(content);
                                    } catch (IOException e) {
                                        sseEmitter.completeWithError(e);
                                    }
                                },
                                error -> {
                                    sseEmitter.completeWithError(error);
                                },
                                () -> {
                                    sseEmitter.complete(); // 在流结束时调用complete()
                                }
                        );
            } catch (Exception e) {
                sseEmitter.completeWithError(e);
            }
        });

        return sseEmitter;
    }

    @GetMapping("/test")
    public SseEmitter test() throws IOException {
        SseEmitter sseEmitter = new SseEmitter(0l);
        for (int i = 0; i < 50 ; i++) {
            sseEmitter.send(i);
        }
        sseEmitter.complete();
        return  sseEmitter;
    }

}
